# Copyright 2020 Makani Technologies LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Convert a STL text file to a C header file for use in the visualizer."""

import sys
import textwrap

import gflags
import gzip

FLAGS = gflags.FLAGS
gflags.DEFINE_string('input_file', None,
                     'Full path to STL model file in text format.',
                     short_name='i')
gflags.DEFINE_string('model_name', None,
                     'Defines the name of the model that is being generated.')
gflags.DEFINE_string('output_file', None,
                     'Full path to output header file.',
                     short_name='o')
gflags.DEFINE_float('scale_factor', 1e-3,
                    'Scale factor for model geometry to convert to meters.')
gflags.DEFINE_bool('gunzip', False, 'Whether the STL model is compressed.',
                   short_name='z')


def GetFaces(stl):
  """Get faces from STL file text."""
  solid_found = False
  facets = []
  for line_num, line in enumerate(stl):
    line = line.strip()
    if line.startswith('solid'):
      if solid_found:
        raise ValueError('Found multiple solids in stl file. (line %d)' %
                         line_num)
      solid_found = True
    elif line.startswith('endsolid'):
      continue
    elif line.startswith('facet'):
      facet = []
    elif line.startswith('endfacet'):
      if len(facet) != 3:
        raise ValueError('Wrong number of facets in component. (line %d)' %
                         line_num)
      facets.append(facet)
    elif line.startswith('outer loop'):
      continue
    elif line.startswith('endloop'):
      continue
    elif line.startswith('vertex'):
      vertex = line[len('vertex'):].strip().split()
      if len(vertex) != 3:
        raise ValueError('Wrong number of vertex components: %s (line %d)' %
                         (line, line_num))
      facet.append(tuple(float(component) * FLAGS.scale_factor
                         for component in vertex))
    else:
      raise ValueError('Unknown line found: %s (line %d)' % (line, line_num))
  if not facets:
    raise ValueError('No facets found in stl file.')
  return facets


def GenerateHeader(stl):
  """Generate C header file."""

  faces = GetFaces(stl)

  model_name = FLAGS.model_name

  if model_name is None:
    assert False, 'gFlag model_name cannot be empty.'
  model_name = model_name.lower().capitalize()

  hdr = textwrap.dedent('''
      // AUTOMATICALLY GENERATED BY vis/generate_model.py

      #ifndef VIS_%sMODEL_H_
      #define VIS_%sMODEL_H_
      // clang-format off

      ''' % (model_name.upper(), model_name.upper()))
  hdr += 'int model%sNumVerts = %d;\n' % (model_name, len(faces) * 3)
  hdr += '\n'
  hdr += 'double model%sVerts [] = {\n' % model_name
  for face_num, face in enumerate(faces):
    hdr += '  // Face %d.\n' % face_num
    for vertex in face:
      hdr += '  %f, %f, %f,\n' % vertex
  hdr += '};\n'
  hdr += textwrap.dedent('''

      // clang-format on
      #endif  // VIS_%sMODEL_H_
      ''' % model_name.upper())
  return hdr


def main(argv):
  """Entry point."""
  try:
    argv = FLAGS(argv)
  except gflags.FlagsError, e:
    print '{}\nUsage: {} ARGS\n{}'.format(e, sys.argv[0], FLAGS)
    sys.exit(1)

  my_open = gzip.open if FLAGS.gunzip else open

  with my_open(FLAGS.input_file, 'r') as input_file:
    with open(FLAGS.output_file, 'w') as output_file:
      output_file.write(GenerateHeader(input_file.readlines()))


if __name__ == '__main__':
  main(sys.argv)
